import torch
from torch import nn


# ----- 定义一个简单的神经网络 ------
class CustomNet(nn.Module):
    def __init__(self):
        super().__init__()
        # 卷积
        self.conv1 = nn.Conv2d(3, 64, 3, 1, 1)
        self.conv2 = nn.Conv2d(64, 128, 3, 1, 1)
        # 池化
        self.pool = nn.MaxPool2d(2, 2)
        # 将原本的矩阵数据展平
        self.flatten = nn.Flatten()
        # 线性回归
        self.fc1 = nn.Linear(128 * 7 * 7, 1024)
        self.fc2 = nn.Linear(1024, 10)  # 10就是图像分类的数量
        # 激活函数
        self.relu = nn.ReLU()
        self.softmax = nn.LogSoftmax(dim=-1)
        # 随机削减神经元（提高准确率）
        self.dropout = nn.Dropout()

    # 前向传播
    def forward(self, x):
        x = self.pool(self.conv1(x))
        x = self.pool(self.conv2(x))
        x = self.flatten(x)
        x = self.dropout(x)
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.relu(self.fc2(x))
        return self.softmax(x)


image = torch.randn(1, 3, 28, 28)
model = CustomNet()
result = model(image)
print(result.shape)  # 输出结果的形状
